import math
import numpy as np
from collections import deque
import time


class PoseClassifier:
    def __init__(self):
        self.pose_history = deque(maxlen=10)
        self.angle_thresholds = {
            "stand_hip": 160,  # 站立时髋关节角度
            "stand_knee": 160,  # 站立时膝关节角度
            "squat_hip": 90,  # 深蹲时髋关节角度
            "squat_knee": 90,  # 深蹲时膝关节角度
            "bow_angle": 45,  # 鞠躬时上半身倾斜角度
            "hands_up": 160,  # 举手时手臂伸展角度
            "wave_range": 0.05,  # 挥手时手部移动幅度
            "sit_hip": 80,  # 坐姿时髋关节角度
            "side_angle": 30,  # 侧身时肩部角度
            "stability_threshold": 0.02,  # 姿态稳定性阈值
            "lie_angle": 15,  # 躺姿时身体倾斜角度
            "run_movement": 0.1,  # 跑步时身体移动幅度
        }

        # 动作持续时间跟踪
        self.action_start_time = 0
        self.current_action = None
        self.action_duration = 0.5  # 动作需要持续的秒数

        # 手部动作状态跟踪
        self.hand_positions = deque(maxlen=15)  # 跟踪更多帧的手部位置

        # 新增：关键点平滑处理和稳定性控制
        self.landmark_history = {
            i: deque(maxlen=5) for i in range(33)
        }  # 存储每个关键点的历史位置
        self.stability_threshold = 0.005  # 稳定性阈值
        self.smooth_factor = 0.7  # 平滑因子

        # 新增：身体位置历史记录
        self.body_positions = deque(maxlen=10)

    def calculate_angle(self, p1, p2, p3):
        """计算三个点形成的角度"""
        a = np.array([p1.x, p1.y])
        b = np.array([p2.x, p2.y])
        c = np.array([p3.x, p3.y])

        ba = a - b
        bc = c - b

        cosine_angle = np.dot(ba, bc) / (np.linalg.norm(ba) * np.linalg.norm(bc))
        angle = np.arccos(cosine_angle)

        return np.degrees(angle)

    def get_body_angle(self, landmarks):
        """计算身体倾斜角度"""
        spine_x = landmarks[24].x - landmarks[12].x
        spine_y = landmarks[24].y - landmarks[12].y
        angle = math.degrees(math.atan2(spine_y, spine_x))
        return angle

    def calculate_complex_angles(self, landmarks):
        """计算更多关节角度"""
        angles = {
            "left_arm": self.calculate_angle(
                landmarks[11], landmarks[13], landmarks[15]  # 左肩-肘-腕
            ),
            "right_arm": self.calculate_angle(
                landmarks[12], landmarks[14], landmarks[16]  # 右肩-肘-腕
            ),
            "left_leg": self.calculate_angle(
                landmarks[23], landmarks[25], landmarks[27]  # 左髋-膝-踝
            ),
            "right_leg": self.calculate_angle(
                landmarks[24], landmarks[26], landmarks[28]  # 右髋-膝-踝
            ),
            "left_knee": self.calculate_angle(
                landmarks[23], landmarks[25], landmarks[27]  # 左髋-左膝-左踝
            ),
            "right_knee": self.calculate_angle(
                landmarks[24], landmarks[26], landmarks[28]  # 右髋-右膝-右踝
            ),
            "spine": self.calculate_angle(
                landmarks[0], landmarks[23], landmarks[24]  # 头-左髋-右髋
            ),
            "shoulders": self.calculate_angle(
                landmarks[11], landmarks[12], landmarks[24]  # 左肩-右肩-右髋
            ),
        }
        return angles

    def detect_movement(self, landmarks):
        """检测动态动作"""
        current_y = landmarks[0].y  # 使用头部位置作为参考
        self.prev_y_positions.append(current_y)

        if len(self.prev_y_positions) >= 10:
            y_variance = np.var(list(self.prev_y_positions))
            y_diff = max(self.prev_y_positions) - min(self.prev_y_positions)

            if y_diff > self.angle_thresholds["jump_height"]:
                return "跳跃", y_diff
            elif y_variance > 0.001:  # 检测剧烈运动
                return "运动中", y_variance
        return None, 0

    def calculate_pose_score(self, angles, target_pose):
        """计算姿势标准度评分"""
        score = 0
        if target_pose == "站立":
            score = 100 - abs(180 - angles["left_leg"]) - abs(180 - angles["right_leg"])
        elif target_pose == "深蹲":
            score = 100 - abs(90 - angles["left_leg"]) - abs(90 - angles["right_leg"])
        return max(0, min(100, score))

    def check_hands_up(self, landmarks):
        """检查双手举过头顶"""
        # 检查双手腕位置是否高于头部
        head_y = landmarks[0].y  # 鼻尖位置
        left_wrist_y = landmarks[15].y
        right_wrist_y = landmarks[16].y

        # 计算手臂伸展角度
        left_arm_angle = self.calculate_angle(
            landmarks[11], landmarks[13], landmarks[15]  # 左肩  # 左肘  # 左腕
        )
        right_arm_angle = self.calculate_angle(
            landmarks[12], landmarks[14], landmarks[16]  # 右肩  # 右肘  # 右腕
        )

        # 手腕需要高于头部，且手臂需要充分伸展
        is_hands_up = (
            left_wrist_y < head_y
            and right_wrist_y < head_y
            and left_arm_angle > self.angle_thresholds["hands_up"]
            and right_arm_angle > self.angle_thresholds["hands_up"]
        )

        confidence = min(landmarks[15].visibility, landmarks[16].visibility)
        return is_hands_up, confidence

    def check_waving(self, landmarks):
        """检查挥手动作"""
        # 记录当前手腕位置
        left_wrist = np.array([landmarks[15].x, landmarks[15].y])
        right_wrist = np.array([landmarks[16].x, landmarks[16].y])
        self.hand_positions.append((left_wrist, right_wrist))

        if len(self.hand_positions) < 10:  # 需要足够的历史数据
            return False, 0.0

        # 计算手腕移动轨迹的方差
        left_trajectory = np.array([pos[0] for pos in self.hand_positions])
        right_trajectory = np.array([pos[1] for pos in self.hand_positions])

        # 分别计算左右手的移动幅度
        left_movement = np.var(left_trajectory, axis=0)
        right_movement = np.var(right_trajectory, axis=0)

        # 判断是否有足够的水平或垂直移动
        is_waving = np.max(left_movement) > 0.005 or np.max(right_movement) > 0.005

        # 计算置信度
        confidence = min(landmarks[15].visibility, landmarks[16].visibility)
        return is_waving, confidence

    def calculate_joint_stability(self, landmarks, joint_indices):
        """计算关节点的稳定性"""
        stability_scores = []
        for idx in joint_indices:
            if len(self.pose_history) >= 5:
                positions = [pose[idx] for pose in self.pose_history]
                variance = np.var([p.y for p in positions])
                stability_scores.append(
                    1.0
                    - min(variance / self.angle_thresholds["stability_threshold"], 1.0)
                )
        return np.mean(stability_scores) if stability_scores else 0.5

    def check_standing(self, angles, landmarks):
        """检查站立姿势"""
        try:
            hip_angle = min(angles["left_leg"], angles["right_leg"])
            knee_angle = min(
                angles["left_leg"], angles["right_leg"]
            )  # 修改这里，使用 leg 而不是 knee

            # 检查垂直度
            spine_vertical = abs(90 - abs(angles["spine"])) < 20

            # 检查双脚位置
            feet_level = abs(landmarks[27].y - landmarks[28].y) < 0.1

            is_standing = (
                hip_angle > self.angle_thresholds["stand_hip"]
                and knee_angle > self.angle_thresholds["stand_knee"]
                and spine_vertical
                and feet_level
            )

            confidence = min(landmarks[23].visibility, landmarks[24].visibility)
            stability = self.calculate_joint_stability(landmarks, [23, 24, 25, 26])

            return is_standing, confidence * stability
        except KeyError as e:
            print(f"姿态检测错误: {e}")
            return False, 0.0

    def check_squatting(self, angles, landmarks):
        """检查深蹲姿势"""
        try:
            hip_angle = min(angles["left_leg"], angles["right_leg"])
            knee_angle = min(angles["left_leg"], angles["right_leg"])  # 同样修改这里

            # 检查重心是否居中
            center_balanced = abs(landmarks[23].x - landmarks[24].x) < 0.1

            # 检查脚跟是否着地
            heels_down = True  # 简化这个检查，因为脚踝点更可靠

            is_squatting = (
                hip_angle < self.angle_thresholds["squat_hip"]
                and knee_angle < self.angle_thresholds["squat_knee"]
                and center_balanced
                and heels_down
            )

            confidence = min(landmarks[23].visibility, landmarks[24].visibility)
            stability = self.calculate_joint_stability(landmarks, [23, 24, 25, 26])

            return is_squatting, confidence * stability
        except KeyError as e:
            print(f"姿态检测错误: {e}")
            return False, 0.0

    def smooth_landmarks(self, landmarks):
        """平滑关键点位置"""
        smoothed_landmarks = []

        for i, landmark in enumerate(landmarks):
            # 将当前关键点添加到历史记录
            self.landmark_history[i].append((landmark.x, landmark.y, landmark.z))

            if len(self.landmark_history[i]) >= 3:
                # 计算平滑后的位置
                x = sum(p[0] for p in self.landmark_history[i]) / len(
                    self.landmark_history[i]
                )
                y = sum(p[1] for p in self.landmark_history[i]) / len(
                    self.landmark_history[i]
                )
                z = sum(p[2] for p in self.landmark_history[i]) / len(
                    self.landmark_history[i]
                )

                # 应用平滑因子
                landmark.x = (
                    landmark.x * (1 - self.smooth_factor) + x * self.smooth_factor
                )
                landmark.y = (
                    landmark.y * (1 - self.smooth_factor) + y * self.smooth_factor
                )
                landmark.z = (
                    landmark.z * (1 - self.smooth_factor) + z * self.smooth_factor
                )

            smoothed_landmarks.append(landmark)

        return smoothed_landmarks

    def check_stability(self, landmarks):
        """检查姿态是否稳定"""
        if not all(len(history) >= 3 for history in self.landmark_history.values()):
            return False

        for i, landmark in enumerate(landmarks):
            history = self.landmark_history[i]
            variance = np.var([p[1] for p in history])  # 主要检查y轴方向的抖动
            if variance > self.stability_threshold:
                return False
        return True

    def check_lie(self, landmarks):
        """检查躺姿"""
        body_angle = abs(self.get_body_angle(landmarks))
        is_lie = body_angle < self.angle_thresholds["lie_angle"]
        confidence = min(
            landmarks[0].visibility, landmarks[23].visibility, landmarks[24].visibility
        )
        return is_lie, confidence

    def check_running(self, landmarks):
        """检查跑步姿势"""
        current_position = np.array([landmarks[0].x, landmarks[0].y])
        self.body_positions.append(current_position)

        if len(self.body_positions) >= 10:
            position_diff = np.max(np.array(self.body_positions), axis=0) - np.min(
                np.array(self.body_positions), axis=0
            )
            movement = np.linalg.norm(position_diff)
            is_running = movement > self.angle_thresholds["run_movement"]
            confidence = min(
                landmarks[0].visibility,
                landmarks[23].visibility,
                landmarks[24].visibility,
            )
            return is_running, confidence
        return False, 0.0

    def classify_pose(self, results):
        if not results.pose_landmarks:
            return "未检测到人体", 0.0

        # 应用平滑处理
        landmarks = self.smooth_landmarks(results.pose_landmarks.landmark)

        # 检查姿态稳定性
        is_stable = self.check_stability(landmarks)
        if not is_stable:
            return self.pose_history[-1] if self.pose_history else ("姿态不稳定", 0.5)

        angles = self.calculate_complex_angles(landmarks)
        current_time = time.time()

        # 检查各种姿势
        pose_checks = [
            (self.check_standing(angles, landmarks), "站立"),
            (self.check_squatting(angles, landmarks), "深蹲"),
            (self.check_hands_up(landmarks), "双手上举"),
            (self.check_waving(landmarks), "挥手"),
            (self.check_lie(landmarks), "躺姿"),
            (self.check_running(landmarks), "跑步"),
        ]

        # 找出置信度最高的姿势
        valid_poses = [
            (pose_name, conf) for (is_pose, conf), pose_name in pose_checks if is_pose
        ]

        if valid_poses:
            best_pose = max(valid_poses, key=lambda x: x[1])
            pose_name, confidence = best_pose

            # 计算姿势评分
            score = self.calculate_pose_score(angles, pose_name)

            # 添加姿势稳定性评分
            stability = self.calculate_joint_stability(
                landmarks, [11, 12, 23, 24, 25, 26]
            )
            final_score = (score + stability * 100) / 2

            return f"{pose_name} (评分: {final_score:.0f})", confidence

        # 检查其他特殊姿势
        spine_angle = abs(angles["spine"])
        if spine_angle > 45:  # 鞠躬判断
            return "鞠躬", 0.8
        elif abs(angles["shoulders"] - 90) < 20:  # 侧身判断
            return "侧身", 0.75
        elif all(
            angle < self.angle_thresholds["sit_hip"]
            for angle in [angles["left_leg"], angles["right_leg"]]
        ):
            return "坐姿", 0.85

        return "其他姿势", 0.5
